# -*- coding:utf-8 -*-
import numpy as np
import torch
import torch.nn.functional as F
import torch.nn as nn

from torchvision.models import resnet18, resnet34, resnet50, resnet152

# 针对单通道
# from efficient_capsule_model_oil import ConvEncoder, EfficientCapsNet, EfficientCapsNetWithReconstruction
from model.efficient_capsnet_orgin_oil_gray import EfficientCapsNet, EfficientCapsNetWithReconstruction
from model.efficient_capsnet_orgin import EfficientCapsOrginNet
from model.efficient_resnet18_capsule_model import EfficientResNet18CapsNet
# 针对三通道
from model.efficient_res_capsule_model import ConvEncoder, EfficientResCapsNet, EfficientResCapsNetWithReconstruction
from model.bilateralfsunet import BilateralFSUnet
from model.simple_efficient_res_capsule_model import SimpleEfficientResCapsNet
from resizer_efficient_resnet18_capsule_model import ResizerEfficientResNet18CapsNet
from model.glom_pytorch import Glom
from crate_encoder import *
from torchsummary import summary

def getmodel(arch, user_device):

    # backbone = resnet18()

    if arch == "resnet18-cifar":
        backbone = resnet18()
        backbone.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
        backbone.maxpool = nn.Identity()
        backbone.fc = nn.Identity()
        return backbone, 512
    elif arch == "resnet18-imagenet":
        backbone = resnet18()
        backbone.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        backbone.fc = nn.Identity()
        return backbone, 512
    elif arch == "resnet18-tinyimagenet":
        backbone = resnet18()
        backbone.avgpool = nn.AdaptiveAvgPool2d(1)
        backbone.fc = nn.Identity()
        return backbone, 512
    elif arch == "resnet152":
        backbone = resnet152()
        backbone.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        backbone.fc = nn.Identity()
        return backbone, 2048
    elif arch == "resnet50":
        backbone = resnet50()
        backbone.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)
        backbone.fc = nn.Identity()
        return backbone, 2048
    elif arch == "conv-encoder":
        backbone = ConvEncoder()
        backbone.fc = nn.Identity()
        return backbone, 256
    elif arch == "efficient-capsule-orgin":
        backbone = EfficientCapsOrginNet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 10, 16
    elif arch == "efficient-res-capsule":
        backbone = EfficientResCapsNet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 256, 16
    elif arch == "efficient-resnet18-capsule":
        backbone = EfficientResNet18CapsNet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 64, 16
    elif arch == 'resizer_efficient-resnet18-capsule':
        backbone = ResizerEfficientResNet18CapsNet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 256, 16
    elif arch == "efficient-capsule":
        backbone = EfficientCapsNet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 512, 16
    elif arch == "simple_efficient_res_capsule":
        backbone = SimpleEfficientResCapsNet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 288, 16
    elif arch == "bilateralfsunet":
        backbone = BilateralFSUnet(user_device)
        backbone.fc = nn.Identity()
        return backbone, 128
    elif arch == 'crate_tiny':
        backbone = CRATE_tiny()
        backbone.fc = nn.Identity()
        return backbone, 384
    elif arch == "crate_small":
        backbone = CRATE_small()
        backbone.fc = nn.Identity()
        return backbone, 576
    elif arch == "crate_base":
        backbone = CRATE_base()
        backbone.fc = nn.Identity()
        return backbone, 768
    elif arch == "crate_large":
        backbone = CRATE_large()
        backbone.fc = nn.Identity()
        return backbone, 1024
    elif arch == "crate_tiny_small":
        backbone = CRATE_tiny_small()
        backbone.fc = nn.Identity()
        return backbone, 128
    elif arch == "glom":
        backbone = Glom(
            dim=64,  # dimension
            levels=4,  # number of levels
            image_size=250,  # image size
            patch_size=50,  # patch size
            consensus_self=False,
            local_consensus_radius=0,
            image_chanel=1,
            return_state=2,
            device=user_device
        )
        return backbone, 25, 64
    else:
        raise NameError("{} not found in network architecture".format(arch))


capsule_network_family_list = ['efficient-capsule-orgin', 'efficient-capsule', 'efficient-res-capsule',
                               'res-capsule', 'efficient-resnet18-capsule', 'simple_efficient_res_capsule',
                               'resizer_efficient-resnet18-capsule']
crate_network_family_list = ['crate_small' ,'crate_base' ,'crate_large', 'crate_tiny', 'crate__tiny_small']

device_use = "cuda"
arch = "glom"

input_data_shape = (1,250,250)
# input_data_shape = (1,28,28)

if arch in capsule_network_family_list:
    backbone, num_capsules, capsules_feature_dim = getmodel(arch, device_use)

elif arch in crate_network_family_list:
    backbone, feature_dim = getmodel(arch, device_use)
elif arch == "glom":
    backbone, num_capsules, capsules_feature_dim = getmodel(arch, device_use)
else:
    backbone, feature_dim = getmodel(arch, device_use)

model = backbone.to(device_use)
summary(model, input_data_shape)